"""
从yolov5检测结果中提取与人相关的边界框信息，并将其转换为VIA格式的json文件
"""
from via3_tool import Via3Json
import pickle
import csv
from collections import defaultdict
import os
import cv2
import sys

avaMin_dense_proposals_path = r'D:\file\postgrad\experiment\silver_spare_dataset\avaMin_dense_proposals_train.pkl'
json_path = r'D:\file\postgrad\experiment\silver_spare_dataset\videos_labelframes_resize'

# 加载pkl文件
with open(avaMin_dense_proposals_path, 'rb') as f:
    info = pickle.load(f, encoding='iso-8859-1')

# 定义了属性字典
attributes_dict = {'1': dict(aname='bird',
                             type=2,
                             options={'0': 'moving', '1': 'grooming', '2': 'alerting', '3': 'resting', '4': 'foraging', '5': 'flapping'},
                             default_option_id="",
                             anchor_id='FILE1_Z0_XY1')
                   }

# 统计视频帧数
dirname = '' # 当前处理的视频文件夹名
len_x = {} # 用于记录每个视频文件夹中的图片数量，键是视频文件夹的名称，值是已处理的图片数量
num_images = 0 # 当前视频的图片总数

# 遍历yolo检测结果，统计视频帧数
for i in info:
  temp_dirname = i.split(',')[0] # 视频文件名
  if dirname == temp_dirname: # 当前处理的视频文件夹名
    len_x[dirname] = len_x[dirname] + 1 # 记录该视频的图片数量
  else:
    # 进入下一个视频文件
    dirname = temp_dirname
    len_x[dirname] = 1 # 初始化为1，表示第一帧已被处理

# 处理每个视频文件夹
dirname = '' # 记录当前正在处理的视频文件夹名称
for i in info:
  temp_dirname = i.split(',')[0] # 视频文件名
  if dirname == temp_dirname:
    # 正在循环一个视频文件里的东西

    # 图片ID从1开始计算
    image_id = image_id + 1
    files_img_num = int(i.split(',')[1]) # 视频中包含的帧总数，也就是图片总数

    # 处理漏帧问题
    if files_img_num - 1 != image_id: # 判断当前处理的图片编号是否与该视频中的图片总数不一致，如果不一致说明有未处理的帧或者跳帧的情况
      files_dict[str(image_id)] = dict(fname=i.split(',')[0] + '_' + (str(image_id)).zfill(6) + '.jpg', type=2) # 为当前帧创建一个空标注，生成文件名并存储到字典中，文件名格式为视频文件名_帧编号.jpg，其中帧编号用 6 位数的格式表示
      via3.dumpFiles(files_dict) # 生成一个空标注
      if files_img_num - 1 != image_id: # 再次检查是否还存在未处理的帧
        while image_id < files_img_num - 1: # 检查确保每一帧都被处理
          image_id = image_id + 1
          files_dict[str(image_id)] = dict(
            fname=i.split(',')[0] + '_' + (str(image_id)).zfill(6) + '.jpg', type=2) # 为当前帧生成文件名并存入
          via3.dumpFiles(files_dict)
          print("middle loss", image_id, "   ", num_images)
          print("files_img_num-1", files_img_num - 1, " image_id", image_id) # 打印当前处理的帧编号以及视频中的图片总数，用于调试和跟踪进度
          len_x[dirname] = len_x[dirname] + 1
          continue

    # 提取边界框信息并保存
    files_dict[str(image_id)] = dict(fname=i.split(',')[0] + '_' + (str(image_id)).zfill(6) + '.jpg', type=2) # 为当前帧再次生成文件名并更新，记录图片文件名和类型

    # 提取边界框信息
    for vid, result in enumerate(info[i], 1): # 使用for循环，遍历其中的每个检测结果
      xyxy = result
      xyxy[0] = img_W * xyxy[0]
      xyxy[2] = img_W * xyxy[2]
      xyxy[1] = img_H * xyxy[1]
      xyxy[3] = img_H * xyxy[3]
      temp_w = xyxy[2] - xyxy[0]
      temp_h = xyxy[3] - xyxy[1]

      metadata_dict = dict(vid=str(int(i.split(',')[1])),
                           xy=[2, float(xyxy[0]), float(xyxy[1]), float(temp_w), float(temp_h)],
                           av={'1': '0'}) # 为当前帧生成元数据字典，包括帧编号、边界框的位置信息以及其他附加信息

      metadatas_dict['image{}_{}'.format(str(int(i.split(',')[1])), vid)] = metadata_dict # 存入元数据字典，键为帧编号和检测结果编号的组合

    # 保存文件与元数据
    via3.dumpFiles(files_dict)
    via3.dumpMetedatas(metadatas_dict)

    print("OK ", image_id, "   ", num_images) # 打印当前处理的帧编号和图片总数，表示处理完成
    if image_id == num_images: # 如果相等表示所有帧已处理完成
      views_dict = {} # 初始化，存储视图相关的数据
      for i, vid in enumerate(vid_list, 1): # 遍历视频ID列表，为每个vid（视频ID）创建一个默认字典
        views_dict[vid] = defaultdict(list) # 为每个视频id创建一个默认字典
        views_dict[vid]['fid_list'].append(str(i)) # fid_list存储帧ID，帧的编号从1开始
      via3.dumpViews(views_dict) # 保存视图信息
      via3.dempJsonSave() # 保存当前的JSON文件，确保视图数据被正确保存
      print("save")

    print("image_id", image_id, " len_x[dirname]", len_x[dirname], " num_images", num_images) # 打印当前的图片帧，当前视频文件夹的帧数量，总帧数
    if image_id == len_x[dirname] and image_id < num_images: # 表示有漏帧
      while image_id < num_images:
        image_id = image_id + 1
        files_dict[str(image_id)] = dict(fname=i.split(',')[0] + '_' + (str(image_id)).zfill(6) + '.jpg', type=2) # 生成对应的文件名
        via3.dumpFiles(files_dict)
      print("end loss", image_id, "   ", num_images) # 表示存在帧丢失的情况
      views_dict = {} # 初始化字典，保存新的视图信息
      for i, vid in enumerate(vid_list, 1):
        views_dict[vid] = defaultdict(list)
        views_dict[vid]['fid_list'].append(str(i)) # 为每个vid（视频ID）创建默认字典，并将帧ID加入fid_list
      via3.dumpViews(views_dict)
      via3.dempJsonSave()
      print("save")
  else:
    # 进入下一个视频文件
    dirname = temp_dirname
    print("dirname", dirname)

    temp_json_path = os.path.join(json_path, dirname, f'{dirname}_proposal.json') # 为每一个视频文件创建一个via的json文件

    num_images = 0 # 记录视频帧总数
    img_H = 0 # 每帧图片的高度
    img_W = 0 # 每帧图片的宽度
    for root, dirs, files in os.walk(os.path.join(json_path, dirname), topdown=False): # 遍历标注帧的视频文件夹路径
      if "ipynb_checkpoints" in root:
        continue

      for file in files:
        if file.endswith('.jpg'): # 遍历当前文件夹中的所有文件，是否以jpg结尾
          num_images = num_images + 1 # 统计视频文件夹中的帧数
          temp_img_path = os.path.join(json_path, dirname, file)  # 生成帧图片的完整路径
          img = cv2.imread(temp_img_path)  # 读取图片信息
          sp = img.shape  # 获取图片尺寸
          img_H = sp[0]
          img_W = sp[1]
    via3 = Via3Json(temp_json_path, mode='dump') # 处理json文件
    vid_list = list(map(str, range(1, num_images + 1))) # 视频帧列表，并转换为字符串形式
    via3.dumpPrejects(vid_list) # project
    via3.dumpConfigs() # config
    via3.dumpAttributes(attributes_dict) # attribute

    files_dict, metadatas_dict = {}, {} # 保存每帧图片文件的信息，保存标注的元数据（例如边界框信息）
    # 图片ID从1开始计算
    image_id = 1
    files_dict[str(image_id)] = dict(fname=i.split(',')[0] + '_' + (str(image_id)).zfill(6) + '.jpg', type=2) # ID从1开始，生成当前帧的文件名，并将其保存

    for vid, result in enumerate(info[i], 1): # 遍历检测结果
      xyxy = result
      xyxy[0] = img_W * xyxy[0]
      xyxy[2] = img_W * xyxy[2]
      xyxy[1] = img_H * xyxy[1]
      xyxy[3] = img_H * xyxy[3]
      temp_w = xyxy[2] - xyxy[0]
      temp_h = xyxy[3] - xyxy[1]

      metadata_dict = dict(vid=str(int(i.split(',')[1])),
                           xy=[2, float(xyxy[0]), float(xyxy[1]), float(temp_w), float(temp_h)],
                           av={'1': '0'}) # 构建当前帧的元数据信息
      # print(metadata_dict)
      metadatas_dict['image{}_{}'.format(str(int(i.split(',')[1])), vid)] = metadata_dict # 将当前帧的元数据添加到字典中，键值是image_id和vid组合后的字符串

    via3.dumpFiles(files_dict)
    via3.dumpMetedatas(metadatas_dict)